#
# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES.
# Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
#!/bin/bash

# This script will run after being booted into a installer kernel 
# This will setup the disk, grub etc for the actual SONiC to boot from

# NOTE: Replace these flag at build time
declare -r IMAGE_VERSION="{{IMAGE_VERSION}}"
declare -r INSTALLER_PAYLOAD="{{INSTALLER_PAYLOAD}}"
declare -r FILESYSTEM_DOCKERFS="{{FILESYSTEM_DOCKERFS}}"
declare -r DOCKERFS_DIR="{{DOCKERFS_DIR}}"
declare -r FILESYSTEM_SQUASHFS="{{FILESYSTEM_SQUASHFS}}"
declare -r KERNEL_VERSION="{{KERNEL_VERSION}}"
declare -r BF2_GRUB_CFG="{{BF2_GRUB_CFG}}"
declare -r BF3_GRUB_CFG="{{BF3_GRUB_CFG}}"

declare -r image_dir="image-$IMAGE_VERSION"
declare -r demo_volume_revision_label="SONiC-OS-${IMAGE_VERSION}"

declare -r rshimlog=`which bfrshlog 2> /dev/null`
declare -r distro="SONiC"

declare -r device_nvmv=/dev/nvme0n1
declare -r device_emmc=/dev/mmcblk0

pn=$(dmidecode -t 4 | grep "Part Number" | awk '{split($NF,a,"-"); print tolower(a[1])}')
declare -r platform=arm64-nvda_bf-$pn

# The reboot of the DPU from the installer is required only for standalon platform.
# In Smart Switch the reboot is triggered by the host when the DPU notifies it that the installation is finished.
if [[ $(echo $platform | grep 9009d3b600) != "" ]]; then
	declare -r reboot_is_needed=true
	declare -r fw_upgrade_is_needed=false
else
	declare -r reboot_is_needed=false
	declare -r fw_upgrade_is_needed=true
fi

declare -r capsule=/lib/firmware/mellanox/boot/capsule/boot_update2.cap

run_bash_session()
{
	/bin/bash </dev/ttyAMA0 >/dev/ttyAMA0 2>&1
}

rshim_log()
{
	# Write a message to rshim. Rshim can handle only a limited number of bites.
	# It should be used only for critical messages.
	echo "INFO: $*"
	if [ -n "$rshimlog" ]; then
		$rshimlog "INFO: $*"
	fi
}

log()
{
	# Write message to serial concole
	if [ -e /dev/ttyAMA0 ]; then
		echo "INFO: $*" > /dev/ttyAMA0
	fi
}

function_exists()
{
	declare -f -F "$1" > /dev/null
	return $?
}

#
# Check auto configuration passed from boot-fifo
#

declare -r boot_fifo_path="/sys/bus/platform/devices/MLNXBF04:00/bootfifo"
if [ -e "${boot_fifo_path}" ]; then
	cfg_file=$(mktemp)
	# Get 16KB assuming it's big enough to hold the config file.
	dd if=${boot_fifo_path} of=${cfg_file} bs=4096 count=4

	#
	# Check the .xz signature {0xFD, '7', 'z', 'X', 'Z', 0x00} and extract the
	# config file from it. Then start decompression in the background.
	#
	offset=$(strings -a -t d ${cfg_file} | grep -m 1 "7zXZ" | awk '{print $1}')
	if [ -s "${cfg_file}" -a ."${offset}" != ."1" ]; then
		log "Found bf.cfg"
		cat ${cfg_file} | tr -d '\0' > /etc/bf.cfg
	fi
	rm -f $cfg_file
fi

if [ -e /etc/bf.cfg ]; then
	. /etc/bf.cfg
fi

ex() {
    echo "Executing command: $@" > /dev/ttyAMA0

    local rc=0
    $@ 2>&1 > /dev/ttyAMA0
    rc=$?

    if [[ $rc -ne 0 ]]; then
        echo "RC: $rc" > /dev/ttyAMA0
    fi
}

if (lspci -n -d 15b3: | grep -wq 'a2dc'); then
	module=BF3
	if [[ $DHCP_CLASS_ID != "" ]]; then
		DHCP_CLASS_ID="BF3Client"
	fi
else
	module=BF2
	if [[ $DHCP_CLASS_ID != "" ]]; then
		DHCP_CLASS_ID="BF2Client"
	fi
fi

log "$distro installation started on $module module"

default_device_label="SONiC-OS"

device_label=${device_label:-$default_device_label}

if [[ $module == "BF3" ]]; then
	default_device=$device_nvmv
	if [[ $(blkid -L $device_label) =~ ${device_emmc} ]] && [[ -b ${device_emmc}p1 ]]; then
		# Delete EFi boot partition on eMMC device. This is required to migrate to NVME device and boot from NVME correctly
		sfdisk --force ${device_emmc} --delete 1
	fi
else
	default_device=$device_emmc
fi

device=${device:-$default_device}

log "Using $device device and $device_label device label"

# We cannot use wait-for-root as it expects the device to contain a
# known filesystem, which might not be the case here.
while [ ! -b $device ]; do
    log "Waiting for $device to be ready"
    sleep 1
done

# Flash image
bs=512
reserved=34
boot_size_megs=50
mega=$((2**20))
boot_size_bytes=$(($boot_size_megs * $mega))

disk_sectors=`fdisk -l $device | grep "Disk $device:" | awk '{print $7}'`
disk_end=$((disk_sectors - reserved))
boot_start=2048
boot_size=$(($boot_size_bytes/$bs))
root_start=$((2048 + $boot_size))
root_end=$disk_end
root_size=$(($root_end - $root_start + 1))

dd if=/dev/zero of="$device" bs="$bs" count=1
sfdisk -f "$device" << EOF
label: gpt
label-id: A2DF9E70-6329-4679-9C1F-1DAF38AE25AE
device: ${device}
unit: sectors
first-lba: $reserved
last-lba: $disk_end

${device}p1 : start=$boot_start, size=$boot_size, type=C12A7328-F81F-11D2-BA4B-00A0C93EC93B, uuid=CEAEF8AC-B559-4D83-ACB1-A4F45B26E7F0, name="EFI System", bootable
${device}p2 : start=$root_start ,size=$root_size, type=0FC63DAF-8483-4772-8E79-3D69D8477DE4, uuid=F093FF4B-CC26-408F-81F5-FF2DD6AE139F, name=$device_label
EOF
sync

# Refresh partition table
ex blockdev --rereadpt ${device}

if function_exists bfb_pre_install; then
	log "Running bfb_pre_install from bf.cfg"
	bfb_pre_install
fi

# Generate some entropy
ex mke2fs -F ${device}p2
ex mkdosfs ${device}p1 -n "system-boot"
ex mkfs.ext4 -F ${device}p2 -L $device_label
ex fsck.vfat -a ${device}p1
mkdir -p /mnt
ex mount -t ext4 ${device}p2 /mnt
mkdir -p /mnt/boot/efi
ex mount -t vfat ${device}p1 /mnt/boot/efi
log "Extracting SONiC files"

mkdir -p /mnt/$image_dir

# Extract the INSTALLER_PAYLOAD to the $image_dir
export EXTRACT_UNSAFE_SYMLINKS=1
ex unzip -o /debian/$INSTALLER_PAYLOAD -x $FILESYSTEM_DOCKERFS "platform.tar.gz" -d /mnt/$image_dir
mkdir -p /mnt/$image_dir/$DOCKERFS_DIR
unzip -op /debian/$INSTALLER_PAYLOAD "$FILESYSTEM_DOCKERFS" | tar xz --warning=no-timestamp -f - -C /mnt/$image_dir/$DOCKERFS_DIR

mkdir -p /mnt/$image_dir/platform
unzip -op /debian/$INSTALLER_PAYLOAD "platform.tar.gz" | tar xz --warning=no-timestamp -f - -C /mnt/$image_dir/platform

# Copy in the machine.conf file
cat <<EOF > /mnt/machine.conf
onie_arch=arm64
onie_platform=$platform
EOF

chmod a+r /mnt/machine.conf

sync

if [[ $fw_upgrade_is_needed == "true" ]]; then
	sonic_fs_path="/mnt/$image_dir/fs.squashfs"
	sonic_fs_mountpoint="/tmp/$image_dir-fs"

	ex mkdir -p $sonic_fs_mountpoint
	ex mount -t squashfs $sonic_fs_path $sonic_fs_mountpoint

	kernel_mft=$(chroot $sonic_fs_mountpoint dpkg -l | grep kernel-mft-dkms-modules | awk '/^ii/ {print $2}')
	mft_files=$(chroot $sonic_fs_mountpoint dpkg -L $kernel_mft)

	for f in $mft_files; do
		if [[ $sonic_fs_mountpoint/$f != *.ko ]]; then
			continue
		fi

		insmod "$sonic_fs_mountpoint/$f"
	done

	ex mkdir -p /etc/mlnx/

	ex ln -s /mnt/$image_dir/platform/fw/asic/fw-BF3.mfa /etc/mlnx/fw-BF3.mfa

	ex mst start

	if function_exists bfb_pre_fw_install; then
		log "Running bfb_pre_fw_install from bf.cfg"
		bfb_pre_fw_install
	fi

	ex $sonic_fs_mountpoint/usr/bin/mlnx-fw-upgrade.sh -v
	if [[ $? != 0 ]]; then
		log "ERROR: FW update failed"
	fi

	ex umount $sonic_fs_mountpoint
fi

if function_exists bfb_post_sonic_install; then
	log "Running bfb_post_sonic_install from bf.cfg"
	bfb_post_sonic_install
fi

{% if SECURE_UPGRADE_MODE in ['dev', 'prod'] %}
demo_volume_label="SONiC-OS"
log "creating demo_volume_label=$demo_volume_label dir under EFI partition to include all boot related modules"
mkdir -p /mnt/boot/efi/EFI/$demo_volume_label

if [ ! -f /secure-boot/mmaa64.efi ]; then
   echo "ERROR: /secure-boot/mmaa64.efi file does not exist"
   exit 1
fi

if [ ! -f /secure-boot/shimaa64.efi ]; then
   echo "ERROR: /secure-boot/shimaa64.efi file does not exist"
   exit 1
fi

if [ ! -f /secure-boot/grubaa64.efi ]; then
   echo "ERROR: /secure-boot/grubaa64.efi file does not exist"
   exit 1
fi

log "copying signed shim, mm, grub, grub.cfg from /secure-boot to /boot/efi/EFI/$demo_volume_label directory"
MNT_DIR="/mnt/boot/efi/EFI/$demo_volume_label"
cp /secure-boot/shimaa64.efi $MNT_DIR
cp /secure-boot/grubaa64.efi $MNT_DIR
cp /secure-boot/mmaa64.efi $MNT_DIR
{% else %}
log "Installing GRUB"
{% endif %}
# Create a minimal grub.cfg that allows for:
#   - configure the serial console
#   - allows for grub-reboot to work
#   - a menu entry for the DEMO OS

grub_cfg=$(mktemp)

# Modify GRUB_CMDLINE_LINUX from bf.cfg file if required

if [[ $module == "BF3" ]]; then
    DEFAULT_GRUB_CMDLINE_LINUX=$BF3_GRUB_CFG
else
    DEFAULT_GRUB_CMDLINE_LINUX=$BF2_GRUB_CFG
fi

GRUB_CMDLINE_LINUX=${GRUB_CMDLINE_LINUX:-"$DEFAULT_GRUB_CMDLINE_LINUX"}
export GRUB_CMDLINE_LINUX

# Add a menu entry for the SONiC OS
# Note: assume that apparmor is supported in the kernel
demo_grub_entry="$demo_volume_revision_label"

# Find the grub_cfg_root
uuid=$(blkid ${device}p2 | sed -ne 's/.* UUID=\"\([^"]*\)\".*/\1/p')
if [ -z "$uuid" ]; then
	grub_cfg_root=${device}p2
else
	grub_cfg_root=UUID=$uuid
fi

cat <<EOF >> $grub_cfg
{{GRUB_CFG}}
EOF

# Copy the grub.cfg onto the boot-directory as specified in the grub-install
{% if SECURE_UPGRADE_MODE in ['dev', 'prod'] %}
ex cp $grub_cfg /mnt/boot/efi/EFI/$demo_volume_label/grub.cfg
ex mkdir -p /mnt/grub
ex cp $grub_cfg /mnt/grub/grub.cfg
{% else %}
ex mkdir -p /mnt/grub
ex cp $grub_cfg /mnt/grub/grub.cfg

{% endif %}
sync

log "GRUB CFG Updated"

# Update HW-dependant files

umount /mnt/boot/efi
umount /mnt

ex blockdev --rereadpt ${device}

ex fsck.vfat -a ${device}p1
sync

if function_exists bfb_pre_bootmgr_update; then
    log "Running bfb_pre_bootmgr_update from bf.cfg"
    bfb_pre_bootmgr_update
fi

if [ -e ${capsule} ]; then
    log "Update capsule: ${capsule}"
    ex bfrec --capsule ${capsule}
fi

log "Updating EFIBootMgr"

if [ ! -d /sys/firmware/efi/efivars ]; then
	ex mount -t efivarfs none /sys/firmware/efi/efivars
fi

# Cleanup boot partition
ex bfbootmgr --cleanall
ex rm -f /sys/firmware/efi/efivars/Boot*
ex rm -f /sys/firmware/efi/efivars/dump-*

# If any of the following commands fail or the image is not booting after the installation check if efivars driver is available and loaded
if [[ $(which grub-install) != "" ]]; then
	log "Installing grub with grub-install utility"

	ex mount ${device}p2 /mnt/
	ex mount ${device}p1 /mnt/boot/efi/
	ex grub-install ${device}p1 --bootloader-id=$device_label --locale-directory=/mnt/usr/share/locale --efi-directory=/mnt/boot/efi/ --boot-directory=/mnt/
	ex umount /mnt/boot/efi
	ex umount /mnt
else
	log "Updating bootmgr with efibootmgr utility"

	if efibootmgr | grep $device_label; then
		ex efibootmgr --delete-bootnum  -b "$(efibootmgr | grep $device_label | cut -c 5-8)"
	fi
	ex efibootmgr -c -d "$device" -p 1 -L $device_label -l "\EFI\\$device_label\grubaa64.efi"
fi
{% if SECURE_UPGRADE_MODE in ['dev', 'prod'] %}
uefi_part=1
ex efibootmgr --create \
    --label "$demo_volume_label" \
    --disk "$device" --part $uefi_part \
    --loader "\EFI\$demo_volume_label\shimaa64.efi" || {
    echo "ERROR: efibootmgr failed to create new boot variable on: $device"
    exit 1
}
echo "uefi_shim: Secure Boot components installed successfully"
{% endif %}
BFCFG=`which bfcfg 2> /dev/null`
if [ -n "$BFCFG" ]; then
	# Create PXE boot entries
	# Not adding CX ifaces because presumably they'll not be used for PXE
	if [ -e /etc/bf.cfg ]; then
		mv /etc/bf.cfg /etc/bf.cfg.orig
	fi

	cat > /etc/bf.cfg << EOF
BOOT0=DISK
BOOT1=NET-OOB-IPV4
BOOT2=NET-OOB-IPV6
BOOT3=NET-RSHIM-IPV4
BOOT4=NET-RSHIM-IPV6
PXE_DHCP_CLASS_ID=$DHCP_CLASS_ID
EOF

	ex $BFCFG

	# Restore the original bf.cfg
	/bin/rm -f /etc/bf.cfg
	if [ -e /etc/bf.cfg.orig ]; then
		mv /etc/bf.cfg.orig /etc/bf.cfg
	fi
fi

if [ -n "$BFCFG" ]; then
	ex $BFCFG
fi

if function_exists bfb_post_install; then
	ex bfb_post_install
fi


rshim_log "Installation finished"

if [[ $reboot_is_needed == "true" ]]; then
	rshim_log "Rebooting..."
	reboot -f
else
	rshim_log "Waiting for reset from the host..."
	while true; do
		sleep 1
	done
fi
